package com.pan.admin.generator.server;

import com.gitee.apanlh.util.base.MapUtils;
import com.gitee.apanlh.util.cache.local.Cache;
import com.gitee.apanlh.util.cache.local.CacheUtils;
import com.gitee.apanlh.util.log.Log;
import com.pan.admin.generator.util.ApiGenAIStreamUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 *
 *  @author Pan
 */
public class ApiGenWebSocketServer extends TextWebSocketHandler {

    private static final Cache<String, WebSocketSession> cache = CacheUtils.cache(MapUtils.newConcurrentHashMap());

    public ApiGenWebSocketServer() {
        super();
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        cache.put(session.getId(), session);
        Log.get().info("建立成功");
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        WebSocketSession webSocketSession = cache.get(session.getId());
        String clientMessage = message.getPayload();
        Log.get().info("获取消息:{}", clientMessage);

        String sessionId = webSocketSession.getId();
        if (webSocketSession.isOpen()) {
            ApiGenAIStreamUtils.genText(webSocketSession.getId(), clientMessage);

            while (true) {
                boolean status = ApiGenAIStreamUtils.getStatus(sessionId);
                if (status) {
                    webSocketSession.close();
                    break;
                }
                String now = ApiGenAIStreamUtils.getNow(sessionId);
                webSocketSession.sendMessage(new TextMessage(now));
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        Log.get().info("异常处理");
        cache.remove(session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        cache.remove(session.getId());
        Log.get().info("关闭session");
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
